# from huggingface_hub import hf_hub_download
# chkpt_path = hf_hub_download("ybelkada/segment-anything", "checkpoints/sam_vit_b_01ec64.pth"
# from segment_anything import build_sam, SamAutomaticMaskGenerator
# mask_generator = SamAutomaticMaskGenerator(build_sam(checkpoint="/home/zry/experiments/Switch-NeRF/auto-sam/checkpoints/sam_vit_b_01ec64.pth"))
# masks = mask_generator.generate("/home/zry/experiments/Switch-NeRF/001938.jpg")

from PIL import Image
import numpy as np
from segment_anything import SamAutomaticMaskGenerator, sam_model_registry
import matplotlib.pyplot as plt
import random
import warnings
warnings.filterwarnings('ignore', category=UserWarning)
import os

os.environ['PYDEVD_WARN_EVALUATION_TIMEOUT'] = '10.0'  # 设置更长的超时时间
os.environ['PYDEVD_UNBLOCK_THREADS_TIMEOUT'] = '10.0'
os.environ['PYDEVD_THREAD_DUMP_ON_WARN_EVALUATION_TIMEOUT'] = 'True'
os.environ['PYDEVD_INTERRUPT_THREAD_TIMEOUT'] = '5.0'


device = "cuda:0"
# device = "cpu"

sam = sam_model_registry["vit_b"](checkpoint="/home/zry/datasets/codes/auto-sam/checkpoints/sam_vit_b_01ec64.pth").to(device)
mask_generator = SamAutomaticMaskGenerator(sam)

image = Image.open("/building/train/rgbs/000001.jpg")
# image = Image.open("/home/zry/experiments/Switch-NeRF/img.jpg")
image_array = np.array(image)

# mask的数量，根据不同的图片有所不同
masks = mask_generator.generate(image_array)
# print(masks)

for i, mask in enumerate(masks):
    output_image = np.ones_like(image_array) * 255
    mseg = mask['segmentation']
    output_image[mseg] = image_array[mseg]

    image = Image.fromarray(output_image)
    image.save(f"mask_image_%s.png" % i)
    
    x, y, width, height = mask['bbox']
    cropped_image = output_image[y:y+height, x:x+width]
    cimage = Image.fromarray(cropped_image)
    cimage.save(f"cmask_image_%s.png" % i)

# 将所有masks都合在一个图片上
def apply_random_colors_with_blend(image_array, masks, blend_ratio=0.5):

    blend_ratio = max(0, min(1, blend_ratio))
    blended_image = image_array.copy()

    for mask in masks:
        mseg = mask['segmentation']
        random_color = np.array([random.randint(0, 255) for _ in range(3)], dtype=np.uint8)
        original_segment = image_array[mseg]
        blended_color = (blend_ratio * random_color + (1 - blend_ratio) * original_segment).astype(np.uint8)
        blended_image[mseg] = blended_color

    return blended_image

cimg = apply_random_colors_with_blend(image_array, masks)
# plt.savefig(f"cimg.png")  # 感觉这种保存的方式可能会压缩图片
image = Image.fromarray(cimg)
image.save('cimg.png')

combined_image = np.concatenate((image_array, cimg), axis=1)
plt.imshow(combined_image)
plt.axis('off')  # 关闭坐标轴显示
plt.savefig("comparison_image.png", bbox_inches='tight', pad_inches=0)  # 确保你已经绘制了图像。plt.savefig 将保存当前绘制的图像，因此你需要先调用绘图函数（如 plt.plot、plt.imshow 等）绘制图像